#coding=utf-8
import numpy as np
import time
from kinematic_model import freedrop
from controller import kalman_filter

import matplotlib.pyplot as plt
#   支持中文
import matplotlib as mpl
mpl.rcParams['font.family']='SimHei'
plt.rcParams['axes.unicode_minus']=False #用来正常显示负号

class Scene:
    '''
        场景
    '''
    def __init__(self,windSpd=np.array([0.7,0.3,0.0]),\
                 initialSpd=np.array([120.0,0.0,120.0])):
        '''
            我现在是一个防空兵
            防空炮打出一枚炮弹，真实的炮弹轨迹，它可能会受风的影响，可能会有随机因素导致偏离目标导致打不中飞机...
            我们可以使用指挥所观测到的炮弹轨迹，因为炮弹距离很远，所以这个观测不是很靠谱...
            所以我们所使用了卡尔曼滤波算法，得到了一条真实的炮弹轨迹...
        '''
        #   真实的炮弹
        self.realShell=freedrop.FreeDropBinder(windSpd=windSpd,initialSpd=initialSpd)
        #   理论上炮弹的落点
        self.theoShell=freedrop.FreeDropBinder(windSpd=np.array([0.0,0.0,0.0]),initialSpd=initialSpd,randRatio=0.0)
        #   卡卡尔曼滤波器
        self.kf=kalman_filter.KF_Onmi3D()
        self.kf.initState[3:6]=initialSpd

        #   绘图区
        self.fig=plt.figure('炮弹弹道图')
        self.ax = self.fig.gca(projection="3d")


        #   数据缓存
        self.realCoord=[]
        self.theoCoord=[]
        self.kalmanCoord=[]
        self.observeCoord=[]

    def UpdateData(self,delta_t=0.2):
        '''
            更新虚拟环境的数据
        :return:
        '''
        #   真实炮弹轨迹
        self.realShell.StateUpdate(delta_t=delta_t)
        #   理论炮弹轨迹
        self.theoShell.StateUpdate(delta_t=delta_t)
        #   观测到的炮弹轨迹
        self.observeCoord.append(self.realShell.position + np.random.random(3) * self.realShell.position[0]/20.0 - self.realShell.position[0]/40.0)
        #   卡尔曼滤波
        '''
            基于卡尔曼滤波，结合理论炮弹轨迹 对观测的炮弹轨迹进行修正
        '''
        self.kf.Predict(velocity=self.theoShell.spd)
        Hybrid_Position=self.kf.Update(self.observeCoord[-1])
        #   绘图(真实的弹道)
        plt.cla()
        self.ax.set_xlim(0, 1000)
        self.ax.set_ylim(-200, 200)
        self.ax.set_zlim(0, 300)
        self.ax.set_xlabel("X坐标（米）")
        self.ax.set_ylabel("Y坐标（米）")
        self.ax.set_zlabel("X坐标（米）")
        #   计算三个类型的炮弹

        self.realCoord.append(np.copy(self.realShell.position))     #   真实炮弹
        self.theoCoord.append(np.copy(self.theoShell.position))     #   理论模型
        self.kalmanCoord.append(np.copy(Hybrid_Position))


        self.curve2Draw=np.array(self.realCoord)
        self.curve2 = np.array(self.observeCoord)
        self.curve3 = np.array(self.theoCoord)
        self.curve4 = np.array(self.kalmanCoord)

        self.ax.plot(self.curve2Draw[:,0],self.curve2Draw[:,1],self.curve2Draw[:,2],label='真实炮弹',color='red')
        self.ax.scatter(self.curve2[:, 0], self.curve2[:, 1], self.curve2[:, 2],'rv+', label='炮弹观测数据', color='blue',alpha=0.5,s=1)
        self.ax.plot(self.curve3[:, 0], self.curve3[:, 1], self.curve3[:, 2], label='炮弹理论轨迹', color='green', alpha=0.5)
        self.ax.plot(self.curve4[:, 0], self.curve4[:, 1], self.curve4[:, 2], label='炮弹融合轨迹', color='yellow', alpha=1.0)
        self.ax.legend()
        plt.pause(0.05)

#   开始模拟环境
#plt.ion()

s=Scene()

for i in range(1000):
    if s.realShell.position[2]<0: break
    s.UpdateData()
plt.ioff()
plt.show()
